package com.lsa.iview.heaps;

import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class PriorityQueue<E> implements Iterable<E> {
	private static final int INITIAL_SIZE = 24;
	private final Comparator<E> comparator;
	
	private Object[] elements = new Object[INITIAL_SIZE];
	private final Map<Object, Integer> keyIndex = new HashMap<>();
	private int size;

	public PriorityQueue(Comparator<E> comparator) {
		this.comparator = comparator;
	}

	public PriorityQueue(E[] elements, Comparator<E> comparator) {
		this(comparator);
		this.elements =  Arrays.copyOf(elements, elements.length * 2);
		this.size = elements.length;
		for (int i = this.elements.length / 2 - 1; i > -1; i--) {
			heapify(i);
		}
	}
	
	public int size() {
		return size;
	}
	
	public E head() {
		return get(0);
	}
	
	public E extractHead() {
		E e = get(0);
		swap(0, size - 1);
		keyIndex.remove(elements[size - 1]);
		--size;
		heapify(0);
		return e;
	}
	
	public void updateKey(E newKey) {
		Integer index = keyIndex.get(newKey);
		if (index != null) {
			updateKey(index, newKey);
		}
	}
	
	public void updateKey(int i, E newKey) {
		E key = get(i);
		if (comparator.compare(key, newKey) >= 0) {
			return;
		}
		elements[i] = newKey;
		promote(i);
	}

	public void insert(E newKey) {
		size++;
		ensureCapacity();
		elements[size - 1] = newKey;
		keyIndex.put(newKey, size - 1);
		promote(size - 1);
	}
	
	private void ensureCapacity() {
		if (size == elements.length) {
			Object[] newElements = new Object[elements.length * 2];
			for (int i = 0; i < elements.length; i++) {
				newElements[i] = elements[i];
			}
			elements = newElements;
		}
	}
	
	private void promote(int i) {
		while(i > 0 && comparator.compare(get(parentIdx(i)), get(i)) < 0) {
			swap(parentIdx(i), i);
			i = parentIdx(i);
		}
	}

	@SuppressWarnings("unchecked")
	public E get(int i) {
		checkIndex(i);
		return (E)elements[i];
	}

	private void checkIndex(int i) {
		if (i < 0 || i >= size) {
			throw new IndexOutOfBoundsException();
		}
	}
	
	private int parentIdx(int i) {
		return (i+1)/2 - 1;
	}
	
	private int leftIdx(int i) {
		return 2*i+1;
	}
	
	private int rightIdx(int i) {
		return 2*i+2;
	}
	
	private void heapify(int i) {
		int left = leftIdx(i);
		int right = rightIdx(i);
		
		int max = i;
		if (left <= size() - 1 && comparator.compare(get(max), get(left)) < 0) {
			max = left;
		}
		
		if (right <= size() - 1 && comparator.compare(get(max), get(right)) < 0) {
			max = right;
		}
		
		if (max != i) {
			swap(i, max);
			heapify(max);
		}
	}
	
	private void swap(int i, int j) {
		checkIndex(i);
		checkIndex(j);
		Object tmp = elements[i];
		elements[i] = elements[j];
		elements[j] = tmp;
		
		keyIndex.put(elements[i], i);
		keyIndex.put(elements[j], j);
	}

	@Override
	public Iterator<E> iterator() {
		return new Iterator<E>() {
			int i = 0;
			@Override
			public boolean hasNext() {
				return i < size();
			}

			@SuppressWarnings("unchecked")
			@Override
			public E next() {
				return (E)elements[i++];
			}

			@Override
			public void remove() {
				throw new UnsupportedOperationException(); 
			}
		};
	}
}
